package com.wang.tool.dagre;

import com.alibaba.fastjson.JSONObject;
import com.wang.tool.monitor.LinkLogUtil;
import com.wang.tool.util.CommonUtil;

import java.util.*;
import java.util.concurrent.*;
import java.util.function.Function;


/**
 * 图流程执行上下文
 *
 * @Author shaopeng
 */
public class DagreContext<T> {

    /**
     * 开始时间
     */
    private final long beginTime = System.currentTimeMillis();

    /**
     * 引擎名称
     * 非空
     */
    private final String dagreName;

    /**
     * 入参
     * 非空
     */
    private final T param;

    /**
     * 超时时长 ms
     * 为空代表 不限制
     * 达到时间后流程中断,不会进行后续节点
     */
    private final Integer timeout;

    /**
     * 等待结束时间
     * 空代表 无限等待
     */
    private final Long waitTime;

    /**
     * 需要执行的ids
     * 非空
     */
    private final Set<String> runIds;

    /**
     * 节点执行状况
     * 每个节点的执行情况会以nodeId为key存进来
     * 初始化时创建,后续无变更
     */
    private final Map<String, NodeState> nodeStates;

    /**
     * 运行时数据信息
     * 非线程安全 多线程场景请自行考虑是否冲突
     */
    private final JSONObject runInfo = new JSONObject();

    /**
     * 取节点执行情况
     *
     * @param id
     * @return
     */
    public NodeState getNodeState(String id) {
        if (id == null) {
            return null;
        }
        return nodeStates.get(id);
    }

    /**
     * 取节点结果
     *
     * @param id
     * @param type
     * @param <V>
     * @return
     */
    public <V> V getNodeResult(String id, Class<V> type) {
        if (id == null) {
            return null;
        }
        NodeState state = nodeStates.get(id);
        return state == null ? null : (V) state.getResult();
    }


    /**
     * 是否有超时限制
     *
     * @return
     */
    public boolean hasTimeLimit() {
        return timeout != null;
    }

    /**
     * 计算已经执行了多久
     *
     * @return
     */
    public int calRunTime() {
        return (int) (System.currentTimeMillis() - beginTime);
    }

    /**
     * 还剩多久可用时间
     *
     * @return 不限制时长返回null,已超时返回负数
     */
    public Integer calRemainTime() {
        if (waitTime == null) {
            return null;
        }
        return (int) (waitTime - System.currentTimeMillis());
    }


    /**
     * 获取运行时数据
     *
     * @param runKey
     * @param type
     * @param <V>
     * @return
     */
    public <V> V getRunObj(String runKey, Class<V> type) {
        return runKey == null ? null : (V) runInfo.get(runKey);
    }

    /**
     * 获取运行时数据,带默认值
     *
     * @param runKey
     * @param def
     * @param <V>
     * @return
     */
    public <V> V getOrDefRunObj(String runKey, V def) {
        return runKey == null ? null : (V) runInfo.getOrDefault(runKey, def);
    }


    /**
     * 获取运行时数据,带初始逻辑
     *
     * @param runKey
     * @param init
     * @param <V>
     * @return
     */
    public <V> V computeIfAbsentRunObj(String runKey, Function<String, V> init) {
        return runKey == null ? null : (V) runInfo.computeIfAbsent(runKey, init);
    }


    /**
     * 判断节点是否已取消
     *
     * @param nodeId
     * @return
     */
    public boolean isCancel(String nodeId) {
        if (nodeId == null) {
            return false;
        }
        NodeState nodeState = nodeStates.get(nodeId);
        return nodeState != null && nodeState.isCancel();
    }

    /**
     * 尝试取消节点极其上游无需执行的节点
     * 会取消节点及其上游可取消节点
     * 带锁
     *
     * @param nodeId
     * @return 目标节点是否取消成功
     */
    public synchronized boolean tryCancelNode(String nodeId) {
        if (nodeId == null) {
            return false;
        }
        //此次需要运行的节点
        NodeState nodeState = nodeStates.get(nodeId);
        if (nodeState == null) {
            LinkLogUtil.newTagLog(DagreEngine.TAG_TRY_CANCEL_DAGRE_NODE)
                    .v1(nodeId)
                    .v2("notInRunIds")
                    .log();
            return false;
        }
        //已经结束了
        if (!notDoneIds.contains(nodeId)) {
            LinkLogUtil.newTagLog(DagreEngine.TAG_TRY_CANCEL_DAGRE_NODE)
                    .v1(nodeId)
                    .v2("nodeIsDone")
                    .log();
            return false;
        }

        //标记为取消
        nodeState.setCancel(true);
        LinkLogUtil.newTagLog(DagreEngine.TAG_TRY_CANCEL_DAGRE_NODE)
                .v1(nodeId)
                .v2("success")
                .log();
        //异步任务取消
        Future<?> future = workFutures.get(nodeId);
        if (future != null) {
            //在运行
            future.cancel(true);
        }

        DagreEngine.NodeWrapper<T> nodeWrapper = nodeLibrary.get(nodeId);
        Set<String> meDepends = nodeWrapper.getMeDepends();
        for (String meDepend : meDepends) {
            cancelUpNode(meDepend);
        }
        return true;
    }


    public long getBeginTime() {
        return beginTime;
    }

    public Long getWaitTime() {
        return waitTime;
    }

    public T getParam() {
        return param;
    }

    public Integer getTimeout() {
        return timeout;
    }

    public JSONObject getRunInfo() {
        return runInfo;
    }

    /**
     * 获取需要执行的节点列表
     *
     * @return
     */
    public Set<String> getRunIds() {
        return runIds;
    }


    public Map<String, NodeState> getNodeStates() {
        return nodeStates;
    }

    /**
     * 是否已经结束
     */
    public boolean isDone() {
        return doneFuture.isDone();
    }

    public String getDagreName() {
        return dagreName;
    }


    /*-------------------------分割线-------------------------*/

    /**
     * 节点库
     */
    private final Map<String, DagreEngine.NodeWrapper<T>> nodeLibrary;


    /**
     * 结束future
     * 非空
     */
    private final CompletableFuture<DagreResult> doneFuture;

    /**
     * 工作中任务features
     * 结束时统一取消
     * 非空
     */
    private final Map<String, Future<?>> workFutures;

    /**
     * 还未完成的ids
     * 非空
     */
    private final Set<String> notDoneIds;

    /**
     * 节点枷锁
     * 节点还剩几个可执行
     * 当节点依赖的节点中任一执行完毕,进行初始化
     * 非空
     */
    private final Map<String, Integer> shackles;

    /**
     * @param runIds 非空
     */
    DagreContext(String dagreName, T param, Integer timeout, Map<String, DagreEngine.NodeWrapper<T>> nodeLibrary, Set<String> runIds) {
        this.dagreName = dagreName;
        this.param = param;
        this.timeout = timeout;
        this.nodeLibrary = nodeLibrary;
        if (getTimeout() != null) {
            waitTime = beginTime + getTimeout();
        } else {
            waitTime = null;
        }
        this.runIds = Collections.unmodifiableSet(runIds);
        this.notDoneIds = new HashSet<>(runIds);
        this.shackles = new HashMap<>(runIds.size());
        this.doneFuture = new CompletableFuture<>();
        //取消未完成任务
        this.doneFuture.whenComplete((result, t2) -> cancelNotDoneWorks());
        this.workFutures = new HashMap<>(runIds.size());
        //初始化结果占位对象
        Map<String, NodeState> nodeStates = new HashMap<>(runIds.size());
        //初始化states
        for (String id : runIds) {
            NodeState state = new NodeState();
            state.setId(id);
            nodeStates.put(id, state);
        }
        this.nodeStates = Collections.unmodifiableMap(nodeStates);
    }


    /**
     * 增加任务future
     *
     * @return
     */
    void addWorkFuture(String id, Future<?> workFuture) {
        if (workFutures == null || workFuture == null) {
            return;
        }
        if (doneFuture.isDone()) {
            workFuture.cancel(true);
        }
        workFutures.put(id, workFuture);
    }

    /**
     * 取消未完成的任务
     */
    void cancelNotDoneWorks() {
        workFutures.forEach((key, future) -> {
            if (future.isDone()) {
                return;
            }
            LinkLogUtil.newTagLog(DagreEngine.TAG_DAGRE_CANCEL_WORKER)
                    .v1(key)
                    .v2(String.valueOf(calRunTime()))
                    .log();
        });
    }


    /**
     * 取消上游节点
     *
     * @param nodeId
     */
    private void cancelUpNode(String nodeId) {
        NodeState nodeState = nodeStates.get(nodeId);
        //此次需要运行的节点
        if (nodeState == null) {
            LinkLogUtil.newTagLog(DagreEngine.TAG_TRY_CANCEL_UP_DAGRE_NODE)
                    .v1(nodeId)
                    .v2("notInRunIds")
                    .log();
            return;
        }
        //已经结束了
        if (!notDoneIds.contains(nodeId)) {
            LinkLogUtil.newTagLog(DagreEngine.TAG_TRY_CANCEL_UP_DAGRE_NODE)
                    .v1(nodeId)
                    .v2("nodeIsDone")
                    .log();
            return;
        }
        DagreEngine.NodeWrapper<T> nodeWrapper = nodeLibrary.get(nodeId);
        Set<String> dependMes = nodeWrapper.getDependMes();
        //所有此次需要执行的下游
        for (String nextNode : dependMes) {
            NodeState nextStage = nodeStates.get(nextNode);
            //过滤此次不执行的节点
            if (nextNode == null) {
                continue;
            }
            if (!nextStage.isCancel()) {
                LinkLogUtil.newTagLog(DagreEngine.TAG_TRY_CANCEL_UP_DAGRE_NODE)
                        .v1(nodeId)
                        .v2("nextNeedRun")
                        .log();
                return;
            }
        }
        //标记为已取消
        nodeState.setCancel(true);
        LinkLogUtil.newTagLog(DagreEngine.TAG_TRY_CANCEL_UP_DAGRE_NODE)
                .v1(nodeId)
                .v2("success")
                .log();
        //异步任务取消
        Future<?> future = workFutures.get(nodeId);
        if (future != null) {
            //在运行
            future.cancel(true);
        }

        //找上游
        Set<String> upNodes = nodeWrapper.getMeDepends();
        for (String upNode : upNodes) {
            cancelUpNode(upNode);
        }
    }

    /**
     * 发现需要触发下游节点
     *
     * @param doneNodeWrapper 当前完成的节点
     * @return
     */
    synchronized Collection<DagreEngine.NodeWrapper<T>> tryFires(DagreEngine.NodeWrapper<T> doneNodeWrapper) {
        //下一级id
        Set<String> nextIds = doneNodeWrapper.getDependMes();
        if (CommonUtil.isEmpty(nextIds)) {
            return null;
        }
        List<DagreEngine.NodeWrapper<T>> fires = new ArrayList<>();
        for (String nextId : nextIds) {
            //依赖他的这个节点不需要执行
            if (!runIds.contains(nextId)) {
                continue;
            }
            DagreEngine.NodeWrapper<T> nextNodeWrapper = nodeLibrary.get(nextId);
            //取
            Integer currentRemain = shackles.computeIfAbsent(nextId, k -> nextNodeWrapper.getMeDepends().size());
            //算
            int remainNum = currentRemain - 1;
            //放回
            shackles.put(nextId, remainNum);
            //仅当=0时才会执行
            if (remainNum == 0) {
                fires.add(nextNodeWrapper);
            }
        }
        return fires;
    }


    /**
     * 标记结束
     *
     * @param finishNodeId 触发结束的节点
     * @param result
     */
    synchronized void markDagreDone(String finishNodeId, DagreResult result) {
        if (doneFuture.isDone()) {
            return;
        }
        if (doneFuture.complete(result)) {
            result.setFinishNodeId(finishNodeId);
        }
    }

    /**
     * 标记节点结束
     *
     * @param nodeId 非空
     * @return
     */
    synchronized boolean markNodeDone(String nodeId) {
        notDoneIds.remove(nodeId);
        return notDoneIds.isEmpty();
    }

    /**
     * 等待结果
     *
     * @return
     * @throws ExecutionException
     * @throws InterruptedException
     * @throws TimeoutException
     */
    DagreResult waitGetResult() throws ExecutionException, InterruptedException, TimeoutException {
        if (waitTime == null) {
            return doneFuture.get();
        }
        //多等1ms,避开临界时间
        int waitMillis = (int) (waitTime + 1 - System.currentTimeMillis());
        return doneFuture.get(waitMillis, TimeUnit.MILLISECONDS);
    }

}
